from ultralytics import YOLO



def main():

    # 加载模型
    # model = YOLO(r"E:\项目\松山湖公安分局无人机自动巡检项目\车辆检测\车辆识别标签\visdrone\模型训练权重\hugging_face_yolo_visdrone.pt")
    model = YOLO(r"C:\Users\14159\Desktop\其他工具\工具箱\ai\runs\detect\train10\weights\best.pt")

    # 训练模型
    train_results = model.train(
        data=r"F:\data\02\红外异常行人数据集-重清洗\data.yaml",  # 数据集 YAML 路径
        epochs=3,  # 训练轮次
        imgsz=900,  # 训练图像尺寸
        batch= 12,
        classes= [0],
        device="0",  # 运行设备，例如 device=0 或 device=0,1,2,3 或 device=cpu
        multi_scale =True,
        lrf=0.01,
        cos_lr = True,
        hsv_h = 0.5,
        hsv_s = 0.5,
        hsv_v = 0.5,
        degrees = 180,
        translate = 0.25,
        scale = 1,
        shear = 10,
        perspective = 0.0005,
        flipud = 1,
        fliplr = 1,
        bgr =1,
        mosaic = 1,
        # mixup =1,
        # copy_paste = 1,
        # copy_paste_mode ="mixup",
        erasing = 0.4,
    )

    # 评估模型在验证集上的性能
    metrics = model.val()


    # # 在图像上执行对象检测
    # model = YOLO("runs/detect/train4/weights/best.pt")

    # # Run batched inference on a list of images
    # results = model(["test_data/frame_90.jpg"])  # return a list of Results objects

    # # Process results list
    # for i,result in enumerate(results):
    #     boxes = result.boxes  # Boxes object for bounding box outputs
    #     masks = result.masks  # Masks object for segmentation masks outputs
    #     probs = result.probs  # Probs object for classification outputs
    #     obb = result.obb  # Oriented boxes object for OBB outputs
    #     result.show()  # display to screen
    #     result.save(filename=f"result_{i}.jpg")  # save to disk
    # # print(results)

    # 将模型导出为 ONNX 格式
    # path = model.export(format="onnx")  # 返回导出模型的路径

if __name__ == '__main__':
    main()

